package com.cy.embeddingdemo.testAgent;

 import com.alibaba.fastjson2.JSON;
 import com.alibaba.fastjson2.JSONObject;
 import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.data.message.ChatMessageType;
 import dev.langchain4j.data.message.UserMessage;
 import dev.langchain4j.model.chat.ChatLanguageModel;
 import dev.langchain4j.model.chat.request.ChatRequest;
 import dev.langchain4j.model.chat.response.ChatResponse;
 import dev.langchain4j.model.openai.OpenAiChatModel;
 import dev.langchain4j.model.openai.OpenAiChatRequestParameters;
 import jakarta.annotation.Resource;
 import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

 import java.io.BufferedWriter;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;

 import static com.cy.embeddingdemo.util.PDFReader.extractPDFContents;

@RequestMapping("pdf")
@RestController
public class PdfReaderController {


    ChatLanguageModel chatLanguageModel;

    public PdfReaderController() {


        OpenAiChatModel model = OpenAiChatModel.builder()
                .baseUrl("http://192.168.1.30:8081/v1/")
                .apiKey("EMPTY")
                .modelName("QwQ-32B")
                .timeout(Duration.ofSeconds(600)) // 设置超时时间为60秒

                .build();

        this.chatLanguageModel = model;
    }




    @RequestMapping("process")
    public void process() {

        String folder = "C:\\数据\\2023";
        extractPDFContents(folder);

    }



    /**
     * 遍历指定路径下的所有PDF文件，并提取内容存放到列表中
     * @param directoryPath 要遍历的目录路径
     * @return 包含每个PDF文件内容的列表
     */
    public void extractPDFContents(String directoryPath) {

        Integer count = 0;

        File directory = new File(directoryPath);



        // 遍历目录下的所有文件和子目录
        File[] files = directory.listFiles();
        if (files != null) {
            for (File file : files) {
                // 如果是文件且扩展名是.pdf，则提取内容
                if (file.isFile() && file.getName().toLowerCase().endsWith(".pdf")) {
                    extractPDFContent(file.getAbsolutePath(),file.getName());
                    count++;
                    System.out.println("总量：" + count);
                }
                // 如果是目录，则递归遍历
                else if (file.isDirectory()) {
                    extractPDFContents(file.getAbsolutePath());
                }
            }
        }
    }

    /**
     * 提取单个PDF文件的内容
     * @param filePath PDF文件的路径
     * @return PDF文件的内容
     */
    private void extractPDFContent(String filePath,String fileName) {
        try (PDDocument document = PDDocument.load(new File(filePath))) {
            PDFTextStripper stripper = new PDFTextStripper();
            String content = stripper.getText(document);
            String requestPrompt = systemPrompt + "下面的安全报告的内容：" + content;

            // 创建用户消息

//            String aiAnswer = chatLanguageModel.chat(requestPrompt);
            System.out.println("写入json文件");

            // 调用 doChat 方法
            String chat = chatLanguageModel.chat(requestPrompt);// 正确的调用方式
            parseAiAnswer(fileName, chat);
        } catch (IOException e) {
            System.err.println("读取PDF文件时出错: " + filePath);
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        String aianswer = "好的，以下是经过提炼和总结后的最终JSON输出：\n\n### 攻击方式部分（Attack方式）\n\n```json\n{\n  \"AttackSummary\": {\n    \"Attacks\": [\n      {\n        \"AttackName\": \"PlugX 驻留攻击\",\n        \"Prerequisites\": [\n          \"受害者收到 ISO 溢出文件并打开含有指向合法 Office 文档的快捷方式文件\"\n        ],\n        \"AttackMethod\": \"攻击者通过恶意 ISO 文件 (LNK) 缓冲受害者的命令执行链并植入 PlugX 杀毒软件。\",\n        \"AttackEffect\": [\n          \"受害者的系统被滥用编码并加载 PlugX 载荷。\",\n          \"创建了新的注册表运行项以持久化恶意组件\",\n          \"连接到远程命令和控制服务器，实现北向移动\"\n        ],\n        \"CVE\": [\n          \"\"\n        ],\n        \"Devices\": [\n          {\n            \"device_type\": \"操作系统\",\n            \"name\": \"Windows\"\n          },\n          {\n            \"device_type\": \"恶意终端\",\n            \"name\": \"LMIGuardianSvc.exe, LMIGuardianDll.dll\"\n          }\n        ],\n        \"AssociatedTools\": [\n          {\n            \"tool\": \"sc.exe\",\n            \"description\": \"用于执行生成恶意服务并监视的第一阶段。”\n          },\n          {\n            \"tool\": \"Mask追い personalities targeting organizations\",\n            \"construction of malware and \n          }\n        ]\n      }\n    ],\n    \"AdditionalNotes\": {\n      \"Vulnerabilities\": \"[In report not mentioned specific CVE vulnerabilities.]\",\n      \"Context\": \"[Mustang Panda APT group targeting government and public sector organizations in Europe in recent years]\"\n    }\n  }\n}\n```\n\n### 攻击技术部分（攻击技术）\n\n```json\n{\n  \"attack_technology_name\": \"PlugX病毒执行\",\n  \"target_system\": {\n    \"operating_system\": \"Windows\",\n    \"environment\": \"企业级别的系统\"\n  },\n  \"implementation_steps\": [\n    {\n      \"step_number\": 1,\n      \"step_name\": \"初步访问\",\n      \"description\": \"受害者受到一个带有其他诱饵的恶意电子邮件附带一个 ISO 溢出文件、实施一个简单的 DLL 欺骗攻击。\"\n    },\n    {\n      \"step_number\": 2,\n      \"step_name\": \"加载恶意插件\",\n      \"description\": \"恶意插件 LMIGuardianSvc.exe 被用来加载 LMIGuardianDll.dll，这只是一个伪装亭的恶意功能师 implants PasteB y (Ge iki lax lad  + ea Pelxy (n ana ,  R')\n    }\n  ],\n  \"attack_effects\": [\n    {\n      \"effect\": \"LMIGuardianDa.\",\n      \"effect\": \"LMIGuardiansv Comp OS Secu Express ca \"\n    }\n  ],\n  \"cve_list\": \"无提及具体CVE编号。\",\n  \"associated_devices_or_software\": [\n    {\n      \"device_type\": \"操作系统\",\n      \"name\": \"Windows\"\n    },\n    {\n      \"software_type\": \"恶意终端\",\n      \"name\": \"LMIGuardiansvc – WHOSEY\"\n    },\n    {\n      \"tool\": \"sc.exe\",\n      \"description\": \"与攻击活动相关的工具和恶意软件。\"\n    }\n  ]\n}\n```";
        parseAiAnswer("xxxx", aianswer);


    }




    public static void parseAiAnswer(String fileName, String json) {

        try{
            // 按照指定分隔符分割字符串
            String[] split = json.split("```json");

            if (split.length > 1) {
                // 提取攻击技术部分

                Object o1 = JSON.parseObject(split[1].split("```")[0]);
                writeJsonToFile(fileName + "_tech.json", o1);
                Object o2 = JSON.parseObject(split[1].split("```")[0]);
                writeJsonToFile(fileName + "_plan.json", o2);
            }

        }catch (Exception e){

            HashMap<Object, Object> objectObjectHashMap = new HashMap<>();
            objectObjectHashMap.put("json", json);
            objectObjectHashMap.put("fileName", fileName);
            objectObjectHashMap.put("error",e.getMessage());

            writeJsonToFile(fileName + "_error.json",objectObjectHashMap);

        }

    }

    /**
     * 将 JSON 数据写入到指定目录的文件中
     *
     * @param fileName      文件名
     * @param content    要写入的 JSON 对象
     */
    public static void writeJsonToFile(String fileName, Object content) {
        String directoryPath = "output/json"; // 目标目录

        // 检查目录是否存在，如果不存在则创建
        File directory = new File(directoryPath);
        if (!directory.exists()) {
            directory.mkdirs(); // 创建目录
        }

        // 构建完整的文件路径
        String filePath = directoryPath + File.separator + fileName;

        // 将 JSON 对象转换为字符串

        // 写入文件
        try (BufferedWriter writer = new BufferedWriter(new FileWriter(filePath))) {
            writer.write(JSON.toJSONString(content));
            System.out.println("JSON数据已成功写入文件：" + filePath);
        } catch (IOException e) {
            System.err.println("写入文件时发生错误：" + e.getMessage());
            e.printStackTrace(); // 打印完整的堆栈信息
        }
    }





    public static String systemPrompt = """
                    1、将上传文档的这篇安全报告提炼并总结，把攻击技术部分内容和攻击方式部分内容进行区分
                    2、对上传文档的攻击方式部分，通过对攻击方式内容的分析，对攻击方式名称、前提条件、攻击方式、攻击效果、涉及CVE漏洞列表、关联设备、关联APT组织、关联工具等进行总结，以下2个攻击总结为其它报告的总结内容，可以参考如下：
                    攻击名称： DDoS攻击
                    前提条件： 存在官方网站（如sejm.gov.pl）
                    攻击方式： 通过大量请求淹没服务器
                    攻击效果： 服务器无法访问
                    关联CVE： 无
                    关联设备：
                    
                    攻击名称： 钓鱼攻击
                    前提条件： 具备政府工作人员或公众的联系信息（如邮箱）
                    攻击方式： 通过伪造政府域名和创建虚假网站，利用虚假的财务补偿计划等手段，诱骗员工或公众访问
                    攻击效果： 获取用户的凭证信息，可能进一步控制用户终端或进行恶意活动
                    关联CVE： 无
                    关联设备：
                    
                    3、对上传文档的攻击方式部分，输出json格式，两个json示例如下：
                    {
                      "AttackSummary": {
                        "Attacks": [
                          {
                            "AttackName": "DDoS攻击",
                            "Prerequisites": [
                              "存在官方网站（如sejm.gov.pl）"
                            ],
                            "AttackMethod": "通过大量请求淹没目标服务器",
                            "AttackEffect": "导致服务器无法访问",
                            "Consequences": [
                              "破坏目标机构的正常运作",
                              "中断公共服务",
                              "降低政府公信力"
                            ]，
                            "CVE": [
                              "", “” ]，
                          },
                          {
                            "AttackName": "钓鱼攻击",
                            "Prerequisites": [
                              "具备政府工作人员或公众的联系信息（如邮箱）"
                            ],
                            "AttackMethod": "通过伪造政府域名和创建虚假网站，利用虚假的财务补偿计划等手段，诱骗员工或公众访问",
                            "AttackEffect": [
                              "获取用户的凭证信息",
                              "可能进一步控制用户终端或进行恶意活动"
                            ],
                            "Consequences": [
                              "造成个人数据泄露",
                              "引发财务损失",
                              "破坏社会信任"
                            ]，
                            "CVE": [
                              "", “” ]，
                          }
                        ],
                        "AdditionalNotes": {
                          "ulnerabilities": "报告中未提及具体的溢出漏洞或其他技术细节。",
                          "Context": "攻击由俄罗斯关联的黑客组织'GhostW riter'实施，动机包括政治报复和破坏关键基础设施。"
                        }
                      }
                    }
                    或者这个
                    {
                      "AttackSummary": {
                        "Attacks": [
                          {
                            "AttackName": "Hangul EPS漏洞利用攻击",
                            "Prerequisites": [
                              "攻击目标为Hangul软件用户",
                              "攻击者需要一个包含 CVE-2017-8291 漏洞的 HWP 文件用于初始访问"
                            ],
                            "AttackMethod": "攻击者将恶意EPS文件嵌入到HWP文档中，通过诱骗目标打开文件触发漏洞，执行恶意代码以下载并执行恶意软件（如M2RAT）",
                            "AttackEffect": [
                              "窃取目标设备上的敏感信息",
                              "包括个人身份信息和通信记录"
                            ],
                            "Consequences": [
                              "目标用户的个人数据泄露",
                              "可能导致身份盗窃或进一步的恶意活动",
                              "对攻击目标的机构造成潜在损害"
                            ],
                            "CVE": [
                              "CVE-2017-8291"
                            ],
                            "Devices": [
                              "受影响的设备包括移动手机"
                            ]，
                           "APTOrganization": "APT37"
                          }
                        ],
                        "AdditionalNotes": {
                          "Vulnerabilities": "报告中提到了Hangul EPS漏洞 (CVE-2017-8291) 被用于初始访问。",
                          "Context": "攻击由RedEyes组织实施，主要针对人權 activist、記者和北韓脫北者，目的是竊取敏感信息。攻击手法包括利用共享記憶體進行命令控制和资料外洩。"
                        }
                      }
                    }
                    
                    
                    
                    
                    4、对上传文档的攻击技术部分，通过对攻击技术内容的分析，对攻击技术名称、目标系统、实现步骤、攻击效果、涉及CVE漏洞列表、关联设备等进行总结，以下2个攻击总结为其它报告的总结内容，可以参考如下：
                    攻击技术名称：内核驱动禁用防御技术（“POOR TRY”驱动技术）
                    目标系统：Windows操作系统，特别是目标企业的网络域和企业环境
                    实现步骤：1、初始访问：攻击者通过钓鱼邮件或其他方式，投递恶意文档或恶意软件，获得对目标系统的初始访问权限\s
                    	2、恶意驱动的部署：攻击者利用恶意DLL（如“POOR TRY”驱动）加载内核模式的驱动程序（如fgt.sys），这些驱动通过Windows服务控制工具（如sc.exe）部署到目标系统\s
                    		恶意驱动的示例命令：sc create fgt binPath= %TEMP%\\fgt.sys type= kernel sc start fgt
                    	3、禁用安全软件：恶意驱动通过终止目标系统上的抗病毒软件、防火墙或其他安全进程，禁用目标系统的防御机制\s
                    	4、横向移动：攻击者利用工具（如SharpHound）进行域信任产品渗透，进一步扩大攻击范围\s
                    攻击效果：禁用了目标系统的安全防御措施，为后续的恶意活动（如数据窃取、横向移动或勒索软件攻击）创建了有利的环境。
                    	攻击者能够实现持久化（通过恶意驱动和服务）和目标系统的横向移动
                    关联CVE： 文档中未提及具体的CVE漏洞
                    关联设备：无
                    关联工具或软件：
                    攻击名称： 钓鱼攻击
                    前提条件： 具备政府工作人员或公众的联系信息（如邮箱）
                    攻击方式： 通过伪造政府域名和创建虚假网站，利用虚假的财务补偿计划等手段，诱骗员工或公众访问
                    攻击效果： 获取用户的凭证信息，可能进一步控制用户终端或进行恶意活动
                    关联CVE： 无
                    关联设备或软件：Windows；
                    关联工具和组件：	恶意驱动（如fgt.sys）。
                    		恶意DLL（如“POOR TRY”驱动）。
                    		恶意软件：GuLoader、Cobalt Strike、Netsh、Netwire RAT、Ping、PsExec。
                    		渗透测试工具：SharpHound。
                    
                    5、对上传文档的攻击方式部分，输出json格式，示例如下：
                    {
                      "attack_technology_name": "内核驱动禁用防御技术（“POOR TRY”驱动技术）",
                      "target_system": {
                        "operating_system": "Windows",
                        "environment": "企业网络域和企业环境"
                      },
                      "implementation_steps": [
                        {
                          "step_number": 1,
                          "step_name": "初始访问",
                          "description": "攻击者通过钓鱼邮件或其他方式，投递恶意文档或恶意软件，获得对目标系统的初始访问权限。"
                        },
                        {
                          "step_number": 2,
                          "step_name": "恶意驱动部署",
                          "description": "攻击者利用恶意DLL（如“POOR TRY”驱动）加载内核模式的驱动程序（如`fgt.sys`），并通过Windows服务控制工具（如`sc.exe`）部署到目标系统。例如，用命令`sc create fgt binPath= %TEMP%\\fgt.sys type= kernel`创建服务，然后用`sc start fgt`启动服务。"
                        },
                        {
                          "step_number": 3,
                          "step_name": "禁用安全软件",
                          "description": "恶意驱动终止目标系统上的抗病毒软件、防火墙或其他安全进程，禁用防御机制。"
                        },
                        {
                          "step_number": 4,
                          "step_name": "横向移动",
                          "description": "攻击者利用工具如SharpHound进行域信任产品渗透，扩大攻击范围。"
                        }
                      ],
                      "attack_effects": [
                        {
                          "effect": "禁用目标系统的安全防御措施。"
                        },
                        {
                          "effect": "为后续恶意活动（如数据窃取或勒索软件攻击）创建有利环境。"
                        },
                        {
                          "effect": "实现恶意软件的持久化。"
                        },
                        {
                          "effect": "支持横向移动以扩大攻击范围。"
                        }
                      ],
                      "cve_list": "无提及具体CVE编号。",
                      "associated_devices_or_software": [
                        {
                          "device_type": "操作系统",
                          "name": "Windows"
                        },
                        {
                          "software_type": "恶意驱动",
                          "name": "fgt.sys, POOR TRY"
                        },
                        {
                          "tool": "sc.exe",
                          "description": "用于创建和启动恶意驱动服务。"
                        },
                        {
                          "malware": "GuLoader, Cobalt Strike, Netsh, Netwire RAT, Ping, PsExec",
                          "description": "与攻击活动相关的工具和恶意软件。"
                        }
                      ]
                    }
                    6、就上面的攻击方式和攻击技术两方面的内容，输出最终的json
                    
                    """;

}
